import javax.swing.tree.TreeNode;

/**
 * Created with IntelliJ IDEA.
 * Description:
 * User: 23735
 * Date: 2023-03-24
 * Time: 22:49
 */
public class RBTree {
    public RBTreeNode root;
    static class RBTreeNode{
        public RBTreeNode left;
        public RBTreeNode right;
        public RBTreeNode parent;
//        默认颜色就是红色, 如果是黑色的话, 黑色的节点会增加, 那么其他路径也需要增加黑色的节点
        public COLOR color=COLOR.RED;
        public int val;
        public RBTreeNode(int val){
            this.val=val;
        }
    }

    public boolean insert(int val){
//        先插入这个节点
        RBTreeNode node=new RBTreeNode(val);
//        根为空, 那么新插入的这个节点就是根节点
        if(root==null){
            root=node;
//            注意不要忘了将root 的颜色置为 黑色
//            不然插入第一个节点时就会出错
            root.color=COLOR.BLACK;
            return true;
        }
//        根节点不为空
        RBTreeNode parent=null; // 用来记录当前节点的父亲节点
        RBTreeNode cur=root;
        while(cur!=null){
            parent=cur;
            if(cur.val<val){
                cur=cur.right;
            }else if(cur.val==val){
//                出现相等的值, 直接返回　false
                return false;
            }else{
                cur=cur.left;
            }
        }
//        这里不能这么判断,因为最后 cur 为空, 那么可能 parent 的左右均为空
//        if(cur==parent.left){
//            parent.left=node;
//        }else{
//            parent.right=node;
//        }
        if(parent.val<val){
            parent.right=node;
        }else{
            parent.left=node;
        }

        node.parent=parent;
//        插入之后开始调节颜色, 从 parent 开始调
        cur=node;
        parent=cur.parent;
//        cur 的颜色始终是红色的
        while(parent!=null&&parent.color==COLOR.RED){
            RBTreeNode grandparent=parent.parent;
//               情况一:cur=red, parent=red, uncle=red, grandparent=black
//                这种情况, 不用考虑　　cur是 parent的左还是右
//            也不用考虑, parent 是 grandfather 的左还是右
            if (parent == grandparent.left) {
                RBTreeNode uncle=grandparent.right;
                if(uncle!=null&&uncle.color==COLOR.RED) {
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandparent.color = COLOR.RED;
//                    继续向上调整
                    cur = grandparent;
                    parent = cur.parent;
                }else{
//                    uncle 不存在或者uncle 为黑色
                    if (cur == parent.right) {
//               情况三:cur=red, parent=red, grandparent=black,uncle=black||null
//               parent==grandparent.left, cur=parent.right
//                    先处理情况三, 左旋, 并交换 parent 和 cur ,将情况三改为情况二
                        rotateLeft(parent);
                        RBTreeNode temp=parent;
                        parent=cur;
                        cur=temp;
                    }
//               情况二:cur=red, parent=red, grandparent=black,uncle=black||null
//               parent==grandparent.left, cur=parent.left
//               所以要分 cur = parent.left还是 cur = parent.right, 即情况二和情况三
                    rotateRight(grandparent);
                    grandparent.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            } else {
//                parent==grandparent.right
                RBTreeNode uncle=grandparent.left;
                if(uncle!=null&&uncle.color==COLOR.RED) {
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandparent.color = COLOR.RED;
//                    继续向上调整
                    cur = grandparent;
                    parent = cur.parent;
                }else{
                    if (cur == parent.left) {
//               情况三:cur=red, parent=red, grandparent=black,uncle=black||null
//               parent==grandparent.right, cur=parent.left
//               同样将情况三处理为情况二
                        rotateRight(parent);
                        RBTreeNode temp=parent;
                        parent=cur;
                        cur=temp;
                    } else {
//               情况三:cur=red, parent=red, grandparent=black,uncle=black||null
//               parent==grandparent.right, cur=parent.right
                        rotateLeft(grandparent);
                        grandparent.color = COLOR.RED;
                        parent.color = COLOR.BLACK;
                    }
                }
            }
        }

//        不管前面怎么调整, 最后都需要把根节点调为 黑色
        root.color=COLOR.BLACK;
        return true;
    }


    public void rotateRight(RBTreeNode parent){
//        注意不要忘了, 当前 parent 节点也有可能是 其他节点的子节点
        RBTreeNode pParent=parent.parent;
        RBTreeNode subL=parent.left;
        RBTreeNode subLR=subL.right; //左节点的右节点
//        开始旋转, 分三步
//        一: 先把左节点提上来
//        注意不要忘记讨论 parent 原来就是 根的情况
        if(parent==root){
            root=subL;
            subL.parent=null;
        }else{
            if(parent==pParent.left){
                pParent.left=subL;
            }else{
                pParent.right=subL;
            }
            subL.parent=pParent;
        }
//        二: 原来的父亲节点降低
        subL.right=parent;
        parent.parent=subL;
//        三: LR 放在旧的根节点的 左边
        parent.left=subLR;
//        注意 LR 要判空
        if(subLR!= null){
            subLR.parent=parent;
        }
    }


    /**
     * 左单旋, 一定是 右树高了
     * @param parent
     */
    public void rotateLeft(RBTreeNode parent){
//        注意不要忘了, 当前 parent 节点也有可能是 其他节点的子节点
        RBTreeNode pParent=parent.parent;
        RBTreeNode subR=parent.right;
        RBTreeNode subRL=subR.left; //右节点的左节点
//        开始旋转, 分三步
//        一: 先把右节点提上来
//        注意不要忘记讨论 parent 原来就是 根的情况
        if(parent==root){
            root=subR;
            subR.parent=null;
        }else{
            if(parent==pParent.left){
                pParent.left=subR;
            }else{
                pParent.right=subR;
            }
            subR.parent=pParent;
        }
//        二: 原来的父亲节点降低
        subR.left=parent;
        parent.parent=subR;
//        三: RL 放在旧的根节点的 右边
        parent.right=subRL;
//        注意 RL 要判空
        if(subRL!= null){
            subRL.parent=parent;
        }
    }


    //    中序遍历有序, 能说明是 二叉搜索树
    public void inOrder(RBTreeNode root){
        if(root==null){
            return ;
        }
        inOrder(root.left);
        System.out.print(root.val+" ");
        inOrder(root.right);
    }


    /**
     * 检查是否有两个连续的红色节点
     * @return
     */
    public boolean checkNoDoubleRed(RBTreeNode root){
        if(root==null){
            return true;
        }
        if(root.color==COLOR.RED){
            RBTreeNode parent=root.parent;
            if(parent.color == COLOR.RED) {
                System.out.println("违反了性质：连续出现了两个红色的节点");
                return false;
            }
        }
        return checkNoDoubleRed(root.left)&&checkNoDoubleRed(root.right);
    }

    /**
     *
     * @param expect 期望的每条路径上的黑色节点的数目
     * @param pathNum 这条路径上 , 已经遇到的黑色节点的数目
     * @return
     */
    public boolean checkBlackNum(RBTreeNode root, int expect, int pathNum){
        if(root==null){
            return true;
        }
        if(root.color==COLOR.BLACK){
            pathNum++;
        }
        if(root.left==null&&root.right==null){
            if(pathNum != expect) {
                System.out.println("违反了性质：每条路径上黑色的节点个数是不一样的！");
                return false;
            }
        }
        return checkBlackNum(root.left, expect, pathNum)&&
                checkBlackNum(root.right, expect, pathNum);
    }

    public boolean isRBTree(RBTreeNode root){
        if(root==null){
            return true;
        }
//        如果根节点是黑色的，　那么直接false
        if(root.color!=COLOR.BLACK){
            System.out.println("违反了性质: 根节点不是黑色的!");
            return false;
        }
//        还需要判断每条路径上, 黑色节点的个数是否一样
        RBTreeNode cur=root;
        int count=0;   // 一条路径上应该有的黑色节点的个数
        while(cur!=null){
            if(cur.color==COLOR.BLACK){
                count++;
            }
            cur=cur.left;
        }
//      同时还需要判断 是否出现了 两个红色节点相连的情况
        return checkBlackNum(root, count, 0)&&checkNoDoubleRed(root);
    }

}
